# 根据25c 25d对拉格朗日乘子 和 位移进行求解的测试

# 实现Test1
import time
from dolfinx import mesh, plot, fem, default_real_type, geometry
from dolfinx.fem import functionspace, petsc
from mpi4py import MPI
import pyvista as pv
import numpy as np
import basix
from basix.ufl import element, mixed_element
import ufl
import norm
from petsc4py import PETSc

import tools  # 定义了一些坐标计算相关的函数


def marker(x): return np.isclose(x[0], 1) & np.isclose(x[1], 1)


def true_velocity_function(x):
    """
    速度真解 函数
    """
    u1 = -4 * x[1] * (4 - x[0]**2)**2 * (4 - x[1]**2)
    u2 = 4 * x[0] * (4 - x[0]**2) * (4 - x[1]**2)**2
    return [u1, u2]


def true_velocity_function_ufl(x):
    """
    速度真解 ufl函数
    """
    return ufl.as_vector(true_velocity_function(x))



def lambda_function(x):
    """拉格朗日乘子真解

    Args:
        x (_type_): msh
    """
    return [np.exp(x[0]), np.exp(x[1])]


def lambda_function_ufl(x):
    """拉格朗日乘子真解
        ufl形式
    Args:
        x (_type_): msh
    """
    return ufl.as_vector([ufl.exp(x[0]), ufl.exp(x[1])])


def laplace_u_ufl(x):
    """拉普拉斯算子作用于速度真解
    """
    x_coord = x[0]
    y_coord = x[1]
    
    # 计算u1的拉普拉斯算子（对应SymPy的Laplacian of u1）
    term1 = 24 * y_coord * (x_coord**2 - 4)**2
    term2 = 16 * y_coord * (3 * x_coord**2 - 4) * (y_coord**2 - 4)
    u1 = term1 + term2
    
    # 计算u2的拉普拉斯算子（对应SymPy的Laplacian of u2）
    term3 = -16 * x_coord * (x_coord**2 - 4) * (3 * y_coord**2 - 4)
    term4 = -24 * x_coord * (y_coord**2 - 4)**2
    u2 = term3 + term4
    return ufl.as_vector([u1, u2])


def laplace_u(x):
    # 提取x和y分量（注意：x[0]是x坐标，x[1]是y坐标）
    x_coord = x[0]
    y_coord = x[1]
    
    # 计算u1的拉普拉斯算子（对应SymPy的Laplacian of u1）
    term1 = 24 * y_coord * (x_coord**2 - 4)**2
    term2 = 16 * y_coord * (3 * x_coord**2 - 4) * (y_coord**2 - 4)
    u1 = term1 + term2
    
    # 计算u2的拉普拉斯算子（对应SymPy的Laplacian of u2）
    term3 = -16 * x_coord * (x_coord**2 - 4) * (3 * y_coord**2 - 4)
    term4 = -24 * x_coord * (y_coord**2 - 4)**2
    u2 = term3 + term4
    
    return np.array([u1, u2])


def displacement_function_ufl(x):
    """位移真解
    """
    return true_velocity_function_ufl(x)  # 位移真解和速度真解相同


def displacement_function(x):
    """位移真解
    """
    return true_velocity_function(x)  # 位移真解和速度真解相同


def c(mu, Y):
    """双线性形式 c(μ, Y) 这里ufl.sym()的使用存在困惑 是否有必要?
    """
    A = (ufl.inner(ufl.grad(mu), ufl.grad(Y))) * dx_solid
    B = ufl.inner(mu, Y) * dx_solid
    return A + B



def g(x):
    """固体域体力项 g = -ΔX - λ UFL表达式
        g_partial_B(x) 需要单独通过弱形式处理
    """
    neg_laplace_X = -laplace_u_ufl(x)  # UFL支持直接对向量取负
    neg_lambda = -lambda_function_ufl(x)

    return neg_laplace_X + neg_lambda



def solid_boundary(x):
    """固体边界"""
    return np.isclose(x[0], 1.0) | np.isclose(
        x[0], -1.0) | np.isclose(x[1], -1.0) | np.isclose(x[1], 1.0)


def u_x_derivative_x(x, y):
    # 计算u_x(x, y)对x的导数
    return 16 * x * y * (4 - x**2) * (4 - y**2)


def u_x_derivative_y(x, y):
    # 计算u_x(x, y)对y的导数
    return (12 * y ** 2 - 16) * (4 - x ** 2)**2


def u_y_derivative_x(x, y):
    # 计算u_y(x, y)对x的导数
    return (16 - 12*x**2) * (4 - y**2)**2


def u_y_derivative_y(x, y):
    # 计算u_y(x, y)对y的导数
    return -16 * x * y * (4 - x**2) * (4 - y**2)


def lambda_x_dx(x, y):
    return np.exp(x)

def lambda_x_dy(x, y):
    return np.zeros_like(y)

def lambda_y_dx(x, y):
    return np.zeros_like(x)

def lambda_y_dy(x, y):
    return np.exp(y)


def g_partial_B1(x, y):
    return np.array([-u_x_derivative_y(x, -1), -u_y_derivative_y(x, -1)])


def g_partial_B2(x, y):
    return np.array([u_x_derivative_x(1, y), u_y_derivative_x(1, y)]) 
                    

def g_partial_B3(x, y):
    return np.array([u_x_derivative_y(x, 1), u_y_derivative_y(x, 1)])


def g_partial_B4(x, y):
    return np.array([-u_x_derivative_x(-1, y), -u_y_derivative_x(-1, y)])


def lambda_partial_B1(x, y):
    return np.array([np.zeros_like(x), -np.exp(y)])

def lambda_partial_B2(x, y):
    return np.array([np.exp(x), np.zeros_like(y)])

def lambda_partial_B3(x, y):
    return np.array([np.zeros_like(x), np.exp(y)])

def lambda_partial_B4(x, y):
    return np.array([-np.exp(x), np.zeros_like(y)])


def getStiffnessMatrix(A: PETSc.Mat):
    """
    获取刚度矩阵的值(np.array)
    """
    # 需要导入
    # from petsc4py import PETSc
    m, n = A.getSize()
    matrix_data = np.zeros((m, n))

    # 获取矩阵的所有非零元素
    row, col, val = A.getValuesCSR()

    # 将CSR格式的数据转换为密集矩阵
    for i in range(m):
        for j in range(row[i], row[i + 1]):
            matrix_data[i, col[j]] = val[j]

    return matrix_data

eps = 1e-6  # 正则项系数 用于稳定方程
X_order = 2  # X使用P2元 传入参数=2

# h_solid = 0.25  # 固体网格大小
# h_solid = 0.125  # 固体网格大小
h_solid = 0.0625  # 固体网格大小
# h_solid = 1 / 32
Lx_solid, Ly_solid = 2, 2  # 固体网格长宽
n_solid = [int(Lx_solid / h_solid), int(Ly_solid / h_solid)]
num_solid_cells = 2 * n_solid[0]**2  # 2*N_solid*N_solid


lambda_size = 2 * (n_solid[0] + 1)**2  # 拉格朗日乘子自由度数量
if X_order == 2:
    X_size = 2 * (2 * n_solid[0] + 1)**2
    print("Use P2/P1 element.")
else:
    X_size = 2 * (n_solid[0] + 1)**2  # 固体位移自由度数量
    print("Use P1/P1 element.")

# 打印网格信息
print(f"solid mesh size = {h_solid}, Number of cells = {num_solid_cells}.")

solid_msh = mesh.create_rectangle(
    MPI.COMM_WORLD, [[-1, -1], [1, 1]], n=n_solid, cell_type=mesh.CellType.triangle)

# key为边界边索引 value为边界边的单元索引
solid_boundary_dict = tools.search_boundary_edge_cell(
    solid_msh, counterclock=True)

# 定义有限元空间 https://defelement.org/elements/p1-iso-p2.html
X_p1 = element(
    "Lagrange", solid_msh.basix_cell(), degree=1, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)
X_p2 = element(
    "Lagrange", solid_msh.basix_cell(), degree=2, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)
lambda_p1 = element(
    "Lagrange", solid_msh.basix_cell(), degree=1, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)
P1 = element(
    "Lagrange", solid_msh.basix_cell(), degree=1,  dtype=default_real_type
)  # 定义一个普通的P1元(非向量元) 用于计算基函数值
P2 = element(
    "Lagrange", solid_msh.basix_cell(), degree=2,  dtype=default_real_type
)  # 定义一个普通的P2元(非向量元) 用于计算基函数值


dx_solid = ufl.Measure("dx", domain=solid_msh)
if X_order == 2:
    W = functionspace(solid_msh, X_p2)  # 固体位移函数空间
elif X_order == 1:
    W = functionspace(solid_msh, X_p1)  # 固体位移函数空间  

S = functionspace(solid_msh, lambda_p1)  # 固体拉格朗日乘子函数空间

(X, lambda_) = ufl.TrialFunction(W), ufl.TrialFunction(S)
(Y, mu) = ufl.TestFunction(W), ufl.TestFunction(S)
#


# 设置流体边界自由度
fdim = solid_msh.topology.dim - 1
solid_facets = mesh.locate_entities_boundary(solid_msh, fdim, solid_boundary)

x_solid = ufl.SpatialCoordinate(solid_msh)  # 固体域上的空间坐标


# 刚度矩阵是fenicsx完成的 
a_3_3 = ufl.inner(ufl.grad(X), ufl.grad(Y)) * dx_solid  # a_s(X, Y)
a_3_4 = -c(lambda_, Y)  # c(λ, Y)


a_4_3 = -c(mu, X)  # c(μ, X)
a_4_4 = fem.Constant(solid_msh, eps) * ufl.inner(lambda_, mu) * dx_solid

# a只求解固体域
a_solid = fem.form(
    [
        [a_3_3, a_3_4],
        [a_4_3, a_4_4],
    ]
)


# 求解前准备数据
num_basis_func_X = X_p1.degree * 3
num_basis_func_lambda = lambda_p1.degree * 3

G = tools.GaussIntegrate(order=9)
bb_tree_s = geometry.bb_tree(solid_msh, 2)

solid_msh.topology.create_connectivity(2, 0)  # 创建固体域顶点和边的连接
lambda_dofmap = S.dofmap  # 拉格朗日乘子的自由度映射 因为计算的是c(lambda, v)

X_dofmap = W.dofmap  # 固体位移的自由度映射

S_coords = S.tabulate_dof_coordinates()[:, :-1]  # 拉格朗日乘子自由度坐标(P1元 固体单元)
X_coords = W.tabulate_dof_coordinates()[:, :-1]  # 固体位移自由度坐标(P1元 固体单元)


# 基于fenicsx组装的load_vec
# 定义四条边界  
boundary1 = mesh.locate_entities_boundary(
    solid_msh, solid_msh.topology.dim - 1, lambda x: np.isclose(x[1], -1.0)
)
boundary2 = mesh.locate_entities_boundary(
    solid_msh, solid_msh.topology.dim - 1, lambda x: np.isclose(x[0], 1.0)
)
boundary3 = mesh.locate_entities_boundary(
    solid_msh, solid_msh.topology.dim - 1, lambda x: np.isclose(x[1], 1.0)
)
boundary4 = mesh.locate_entities_boundary(
    solid_msh, solid_msh.topology.dim - 1, lambda x: np.isclose(x[0], -1.0)
)
# 固体网格（假设已定义 solid_msh）

# 合并实体索引和对应的标记（标记需唯一，如1-4）
entities = np.hstack([boundary1, boundary2, 
                      boundary3, boundary4])
tags = np.hstack([np.full_like(boundary1, 1), 
                  np.full_like(boundary2, 2),
                  np.full_like(boundary3, 3),
                  np.full_like(boundary4, 4)])

# 创建MeshTags（关键：将实体与标记关联）
boundary_tags = mesh.meshtags(solid_msh, fdim, entities, tags)
ds_solid = ufl.Measure("ds", domain=solid_msh, subdomain_data=boundary_tags)  # 边界积分
# 设置积分
g1 = fem.Function(W)
g1.interpolate(lambda x: g_partial_B1(x[0], x[1]))

g2 = fem.Function(W)
g2.interpolate(lambda x: g_partial_B2(x[0], x[1]))

g3 = fem.Function(W)
g3.interpolate(lambda x: g_partial_B3(x[0], x[1]))

g4 = fem.Function(W)
g4.interpolate(lambda x: g_partial_B4(x[0], x[1]))

gggg = ufl.inner(g1, Y) * ds_solid(1) + \
    ufl.inner(g2, Y) * ds_solid(2) + \
    ufl.inner(g3, Y) * ds_solid(3) + \
    ufl.inner(g4, Y) * ds_solid(4)

l1 = fem.Function(S)
l2 = fem.Function(S)
l3 = fem.Function(S)
l4 = fem.Function(S)

l1.interpolate(lambda x: lambda_partial_B1(x[0], x[1]))
l2.interpolate(lambda x: lambda_partial_B2(x[0], x[1]))
l3.interpolate(lambda x: lambda_partial_B3(x[0], x[1]))
l4.interpolate(lambda x: lambda_partial_B4(x[0], x[1]))

llll = ufl.inner(l1, Y) * ds_solid(1) + \
    ufl.inner(l2, Y) * ds_solid(2) + \
    ufl.inner(l3, Y) * ds_solid(3) + \
    ufl.inner(l4, Y) * ds_solid(4)
    
    
load_vec = fem.form(
    [
        # 注意laplace(lambda) = lambda
        ufl.inner(g(x_solid), Y) * dx_solid  + ufl.inner(lambda_function_ufl(x_solid), Y) * dx_solid + gggg - llll,
        -c(mu, displacement_function_ufl(x_solid))
    ]
)

# 位移需要固定一个点 否则算不好H1范数
# def marker(x): return np.isclose(x[0], 1.0) 
# X_corner = fem.locate_dofs_geometrical(W, marker=marker)
# X_bc = fem.Function(W)  # 设置边界函数
# X_bc.interpolate(true_velocity_function)

# bc_X = fem.dirichletbc(X_bc, X_corner)  # 设置边界条件
# A = petsc.assemble_matrix_block(a_solid, bcs=[bc_X]); A.assemble()
# b = petsc.assemble_vector_block(load_vec, a_solid, bcs=[bc_X])
# fem.apply_lifting(b.array, a_solid, bcs=[bc_X])

# 手搓
A = petsc.assemble_matrix_block(a_solid); A.assemble()
b = petsc.assemble_vector_block(load_vec, a_solid)
b_arr = b.array
A_np = getStiffnessMatrix(A)  # np.array

# 重新计算 设置投影
if X_order == 1:
    ...
    # A_XY = np.zeros((2 * X_size, 2 * X_size))
    # for i in range(num_solid_cells):
    #     # 对单元进行遍历 设置稳定项
    #     X_idx = X_dofmap.cell_dofs(i)
    #     cell_points = X_coords[X_idx]  # 三角形的顶点
    #     bdt = tools.BasisDerivativeTransformer(cell_points)
        
    #     G1 = tools.GaussIntegrate(order=1)
    #     G3 = tools.GaussIntegrate(order=3)
    #     gs_points_phy1 = G1.to_physical_gauss_points(*cell_points)
    #     gs_points_phy3 = G3.to_physical_gauss_points(*cell_points)
        
    #     P1_basis_value1, _, _ = P1.tabulate(1, G1.points)
    #     P1_basis_value3, _, _ = P1.tabulate(1, G3.points)


        
    #     for alpha in range(num_basis_func_X):
    #         for beta in range(num_basis_func_X):
    #             item_1 = P1_basis_value1[:, alpha] * P1_basis_value1[:, beta]
    #             item_3 = P1_basis_value3[:, alpha] * P1_basis_value3[:, beta]
    #             A_XY[2 * X_idx[beta], 2 * X_idx[alpha]] += bdt.detJ * (np.dot(item_3, G3.weights) - np.dot(item_1, G1.weights))
    #             A_XY[2 * X_idx[beta]+1, 2 * X_idx[alpha]+1] += bdt.detJ * (np.dot(item_3, G3.weights) - np.dot(item_1, G1.weights))
    # A_XY[1] = 1
             
    # 把A_XY加到A_np中
    # A_np[:2 * X_size, :2 * X_size] += A_XY

# 手搓一下边界条件 置1法
# fix_idx = 0
# X_fix_x, X_fix_y = W.tabulate_dof_coordinates()[fix_idx][:-1]  # 固定位移点的坐标
# A_np[2 * fix_idx] = 0
# A_np[2 * fix_idx + 1] = 0
# A_np[2 * fix_idx, 2 * fix_idx], A_np[2 * fix_idx+1, 2 * fix_idx+1] = [1, 1]
# b_arr[2 * fix_idx], b_arr[2 * fix_idx+1] = true_velocity_function([X_fix_x, X_fix_y])

x = tools.sparse_solve(A_np, b_arr)



# x = PETSc.Vec().create(solid_msh.comm)
# x.setSizes(A.size[0]); x.setFromOptions()

# ksp = PETSc.KSP().create(solid_msh.comm)
# ksp.setOperators(A)
# # ksp.setType(PETSc.KSP.Type.PREONLY)
# # pc = ksp.getPC(); pc.setType(PETSc.PC.Type.LU)
# pc = ksp.getPC()

# ksp.setFromOptions() 
# ksp.solve(b, x)

# print("KSP reason:", ksp.getConvergedReason(),
#       "iters:", ksp.getIterationNumber(),
#       "residual:", ksp.getResidualNorm())

# # 核对 ||Ax-b||
# r = b.duplicate(); A.mult(x, r); r.axpy(-1.0, b)
# print("||Ax-b||_2 =", r.norm())
# x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

# solve
X_sol = x[:X_size]
lambda_sol = x[X_size:]



# 计算L2误差
X_l2_err = norm.get_L2_err(
    msh=solid_msh,
    dof_map=X_dofmap,
    num_cells=num_solid_cells,
    u_true_func=displacement_function,
    u_num=X_sol,
    func_space=W,
    order=X_order
)
lambda_l2_err = norm.get_L2_err(
    msh=solid_msh,
    dof_map=lambda_dofmap,
    num_cells=num_solid_cells,
    u_true_func=lambda_function,
    u_num=lambda_sol,
    func_space=S,
    order=1
)
print(f"Displacement L2 error: {X_l2_err:.6e}")
print(f"Lagrange multiplier L2 error: {lambda_l2_err:.6e}\n")


# H1
# 优化H1范数的计算 提取速度对应的刚度矩阵项 A = eye(A1, A2)
K = A_np[:X_size, :X_size]
X_H1_mat_err = norm.get_H1_err_m(
    u_true_func=displacement_function,
    u_num=X_sol,
    K=K,
    func_space=W,
)


X_h1_err = norm.get_H1_err(
    msh=solid_msh,
    dof_map=X_dofmap,
    num_cells=num_solid_cells,
    u_true_func_diff_ls=[u_x_derivative_x, u_x_derivative_y, u_y_derivative_x, u_y_derivative_y],
    u_num=X_sol,
    func_space=W,
    order=X_order
)
lambda_h1_err = norm.get_H1_err(
    msh=solid_msh,
    dof_map=lambda_dofmap,
    num_cells=num_solid_cells,
    u_true_func_diff_ls=[lambda_x_dx, lambda_x_dy, lambda_y_dx, lambda_y_dy],
    u_num=lambda_sol,
    func_space=S,
    order=1
)
print(f"Displacement H1 error: {X_h1_err:.6e}")
print(f"Displacement H1 error(mat ver): {X_H1_mat_err:.6e}")
print(f"Lagrange multiplier H1 error: {lambda_h1_err:.6e}\n")

# max err
X_max_err = norm.get_max_node_err(
    u_true_func=displacement_function,
    u_num=X_sol,   
    func_space=W
)
lambda_max_err = norm.get_max_node_err(
    u_true_func=lambda_function,
    u_num=lambda_sol,   
    func_space=S
)
print(f"Displacement max error: {X_max_err:.6e}")
print(f"Lagrange multiplier max error: {lambda_max_err:.6e}")
